from lxml import etree
import tensorflow as tf
import os
import logging
import cv2
import matplotlib.pyplot as plt

flags = tf.app.flags
flags.DEFINE_string('data_dir', '', 'Root directory to raw PASCAL VOC dataset.')
flags.DEFINE_string('output_path', '', 'Path to output TFRecord')
FLAGS = flags.FLAGS

def recursive_parse_xml_to_dict(xml):
  """Recursively parses XML contents to python dict.

  We assume that `object` tags are the only ones that can appear
  multiple times at the same level of a tree.

  Args:
    xml: xml tree obtained by parsing XML file contents using lxml.etree

  Returns:
    Python dictionary holding XML contents.
  """
  if not xml:
    return {xml.tag: xml.text}
  result = {}
  for child in xml:
    child_result = recursive_parse_xml_to_dict(child)
    if child.tag != 'object':
        result[child.tag] = child_result[child.tag]
    else:
        if child.tag not in result:
            result[child.tag] = []
        result[child.tag].append(child_result[child.tag])
  return {xml.tag: result}

def export_img_in_bbox( voc_path , write_path ):
    #voc_path specify all the parameters needed
    # others follow VOC dataset arrangement manner
    xmls_dir = os.path.join( voc_path , "Annotations" )
    if not os.path.exists( xmls_dir ):
        logging.error( "cannot find Annotations directory!\n" )
        return

    xmls_list = os.listdir( xmls_dir )
    xmls_list = [ x for x in xmls_list if x[-4:] == ".xml"]

    if not os.path.exists( write_path ):
        os.mkdir( write_path )

    for idx , xml in enumerate( xmls_list ):
        if idx %100 == 0:
            logging.info( "On image %d of %d" , idx , len( xmls_list ) )

        path = os.path.join( xmls_dir , xml )
        with tf.gfile.GFile( path , 'r' ) as fid:
            xml_str = fid.read()
            
        xml_etree = etree.fromstring( xml_str )
        xml_data = recursive_parse_xml_to_dict( xml_etree )['annotation']

        # crop image
        img_name = xml_data[ 'filename' ]
        img_name = os.path.join( voc_path , "JPEGImages" , img_name )
        if not os.path.exists( img_name ):
            logging.warning( "img: %s doesn't exists" , img_name )
            continue
        img = cv2.imread( img_name )

        if 'object' in xml_data and len( xml_data['object'] ) >= 1:
            for obj_idx , obj in enumerate( xml_data['object'] ):
                xmin = int( obj['bndbox']['xmin'] )
                xmax = int( obj['bndbox']['xmax'] )
                ymin = int( obj['bndbox']['ymin'] )
                ymax = int( obj['bndbox']['ymax'] )
            
                img_crop = img[ ymin:ymax , xmin:xmax , : ]
                img_crop_name = xml_data['filename'][:-4] + \
                        '_' + str(obj_idx) + '.jpg'

                WRITE_PATH = os.path.join( write_path , img_crop_name )
                cv2.imwrite( WRITE_PATH , img_crop )

    return

if __name__ == "__main__":
    data_dir = FLAGS.data_dir
    write_path = FLAGS.output_path

    export_img_in_bbox( data_dir , write_path )
